#
#  GoogleStorage.py
#  gStorageBrowser
#
#  Created by Wang Shao-chuan on 2011/1/1.
#  Copyright (c) 2011 National Taiwan University. All rights reserved.
#

import boto
import os
import tempfile
import shutil

config = boto.config

class GoogleStroageException(Exception):
  pass

class DulplicateBucketException(GoogleStroageException):
  pass

class BucketNotFoundExcpetion(GoogleStroageException):
  pass

class GoogleStorage(object):
  def __init__(self):
    self.buckets = boto.storage_uri("","gs").get_all_buckets()
    self.items = {}
  
  @property
  def bucketNames(self):
    return [b.name for b in self.buckets]
    
  def hasBucket(self, bucket_name):
    return bucket_name in self.bucketNames
    
  def createBucket(self, bucket_name):
    if self.hasBucket(bucket_name):
      raise DulplicateBucketException("Duplication of bucket name.")
    else:
      bucket_uri = boto.storage_uri(bucket_name, "gs")
      # create the bucket, it may raise boto.exception.GSCreateError for name conflict
      bucket = bucket_uri.create_bucket()
      if bucket:
        self.buckets.append(bucket)
      return bucket

  def deleteBucket(self, bucket_name):
    if not self.hasBucket(bucket_name):
      raise BucketNotFoundExcpetion("Cannot find the bucket name.")
    else:
      # delelete the bucket here
      bucket_uri = boto.storage_uri(bucket_name, "gs")
      bucket_uri.delete_bucket()
  
  def itemNames(self, bucket_name):
    return [item.name for item in self.itemsInBucket(bucket_name)]

  def itemsInBucket(self, bucket_name):
    ''' Returns the cached items for the bucket name or directly get the items from the bucket. '''
    if self.items.get(bucket_name, None):
      return self.items[bucket_name]
    else:
      uri = boto.storage_uri(bucket_name, 'gs')
      self.items[bucket_name] = [item for item in uri.get_bucket()]
      return self.items[bucket_name]
      
  def hasItemInBucket(self, item_name, bucket_name):
    ''' Returns True if the item name in the bucket.'''
    item_name = item_name.strip()
    dst_uri = boto.storage_uri(bucket_name, 'gs')
    bucket = dst_uri.get_bucket()
    return bool(bucket.lookup(item_name))

  def uploadFileToBucket(self, filename, bucket_name, item_name=None, cb=None):
    ''' :type filename, bucket_name, item_name: string.
        :param filename: full path of src file name.
    '''
    item_name = item_name.strip()
    
    dst_uri = boto.storage_uri(bucket_name, 'gs')
    if item_name is None:
      item_name = os.path.basename(filename)
    new_dst_uri = dst_uri.clone_replace_name(item_name)
        
    # copy src file to a temp file
    src_fp = open(filename, 'rb')
    tmp_src_fp = tempfile.TemporaryFile()
    shutil.copyfileobj(src_fp, tmp_src_fp)
    src_fp.close()
    
    # default progress call back function
    if cb is None:
      def cb(downloadBytes, totalBytes):
        print downloadBytes, totalBytes
    
    dst_key = new_dst_uri.new_key()
    tmp_src_fp.seek(0)
    
    dst_key.set_contents_from_file(tmp_src_fp, cb=cb)
    tmp_src_fp.close()
    
    
  def downloadFileFromItemName(self, bucket_name, item_name, dst_filename=None, cb=None):
    item_name = item_name.strip()
    
    # defaulut destination folder: ~/Downloads
    if dst_filename is None:
      dst_filename = u"~" +  os.sep + "Downloads" + os.sep + os.path.basename(item_name)
      dst_filename = os.path.expanduser(dst_filename)
    src_uri = boto.storage_uri(bucket_name + "/" + item_name, "gs")
    src_key = src_uri.get_key()

    # default progress callback function  
    if cb is None:
      def cb(downloadBytes, totalBytes):
        print downloadBytes, totalBytes
  
    tmp_fp = tempfile.TemporaryFile()
    try:
      src_key.get_file(tmp_fp, cb=cb)
    except Exception:
      tmp_fp.close()
      raise
    
    dst_fp = open(dst_filename, 'wb')
    shutil.copyfileobj(tmp_fp, dst_fp)
    tmp_fp.close()
    dst_fp.close()


  def deleteItem(self, bucket_name, item_name):
    ''' :return True if deleted, else False. '''
    ret = False
    item_name = item_name.strip()
    del_uri = boto.storage_uri(bucket_name + "/" + item_name, "gs")
    if self.hasItemInBucket(item_name, bucket_name):
      del_uri.get_key().delete()
      ret = True
    else:
      if item_name[-1] != '/':
        item_name = item_name + '/'
      itemNames = [i.name for i in self.itemsInBucket(bucket_name)]
      to_be_del = [n for n in itemNames if n.startswith(item_name)]
      for item in to_be_del: 
        uri = boto.storage_uri(bucket_name + "/" +  item, "gs")
        uri.get_key().delete()
      ret = bool(to_be_del)
    return ret
    
  def hasCredentials(self):
    config_path = os.path.expanduser('~/.boto')
    if not os.path.exists(config_path):
      return False
    if not config.has_section('Credentials'):
      return False
    if not config.get('Credentials', 'gs_secret_access_key'):
      return False
    if not config.get('Credentials', 'gs_access_key_id'):
      return False
    return True
    
  def createCredentials(self, gs_secret_access_key, gs_access_key_id, replace=True):
    config_path = os.path.expanduser('~/.boto')
    section_name = 'Credentials'
    if os.path.exists(config_path) and not replace:
      return False
    if not config.has_section(section_name):
      config.add_section(section_name)
    config.set(section_name, 'gs_secret_access_key', gs_secret_access_key)
    config.set(section_name, 'gs_access_key_id', gs_access_key_id)
    fout = open(config_path, 'w')
    config.write(fout)
    fout.close()
    return True